package com.chief.sink.rocket;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.rocketmq.client.consumer.DefaultLitePullConsumer;
import org.apache.rocketmq.common.consumer.ConsumeFromWhere;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.protocol.heartbeat.MessageModel;

import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;

public class RocketMqSource extends RichParallelSourceFunction<String> {

    DefaultLitePullConsumer litePullConsumer;

    Boolean running = true;

    @Override
    public void open(Configuration parameters) throws Exception {
        litePullConsumer = new DefaultLitePullConsumer("test_cg");
        litePullConsumer.setNamesrvAddr("192.168.10.54:9876;192.168.10.55:9876");
        litePullConsumer.subscribe("inTest", "*");
        litePullConsumer.setConsumeFromWhere(ConsumeFromWhere.CONSUME_FROM_LAST_OFFSET);
        litePullConsumer.setPullBatchSize(1);
        litePullConsumer.setConsumerPullTimeoutMillis(1000);
        litePullConsumer.setMessageModel(MessageModel.CLUSTERING);
        litePullConsumer.start();
    }

    @Override
    public void run(SourceContext<String> ctx) throws Exception {
        while (running) {
            for (MessageExt messageExt : litePullConsumer.poll(5000)) {
                ctx.collect(new String(messageExt.getBody()));
                litePullConsumer.commitSync();
            }
        }
    }

    @Override
    public void close() throws Exception {
        litePullConsumer.shutdown();
    }

    @Override
    public void cancel() {
        running = false;
    }
}
